//
// Created by syshe on 2022/3/22.
//

#include "Object.h"

namespace c10 {
    // should this use the globalContext?  Can it get a context passed in somehow?
    UndefinedTensorImpl::UndefinedTensorImpl()
            : TensorImpl(DispatchKey::Undefined, caffe2::TypeMeta(), c10::nullopt) {
        set_storage_access_should_throw();
    }

    int64_t UndefinedTensorImpl::size(int64_t d) const {
        TORCH_CHECK(false, "size(dim) called on an undefined Tensor");
    }

    int64_t UndefinedTensorImpl::stride(int64_t d) const {
        TORCH_CHECK(false, "stride(dim) called on an undefined Tensor");
    }

#ifdef DEBUG
    bool UndefinedTensorImpl::has_storage() const {
        TORCH_INTERNAL_ASSERT_DEBUG_ONLY(!storage_, "UndefinedTensorImpl assumes that storage_ is never set");
        return false;
    }
#endif

    void UndefinedTensorImpl::set_storage_offset(int64_t) {
        TORCH_CHECK(false, "set_storage_offset() called on an undefined Tensor");
    }

    IntArrayRef UndefinedTensorImpl::strides() const {
        TORCH_CHECK(false, "strides() called on undefined Tensor");
    }

    const char* UndefinedTensorImpl::tensorimpl_type_name() const {
        return "UndefinedTensorImpl";
    }

    UndefinedTensorImpl UndefinedTensorImpl::_singleton;
}

namespace c10 {
    void IValue::visit(const std::function<bool( const IValue &)>& visitor) const {
        if (visitor(*this)) {
            return;
        }
        switch (this->tag) {
            case Tag::Tuple:
            case Tag::GenericList: {
                c10::ArrayRef<IValue> elems;
                if (isTuple()) {
                    elems = this->toTupleRef().elements();
                } else {
                    elems = this->toListRef();
                }
                for (auto &elem : elems) {
                    elem.visit(visitor);
                }
                break;
            }
            case Tag::GenericDict:
                for (const auto &pair : this->toGenericDict()) {
                    pair.value().visit(visitor);
                    pair.key().visit(visitor);
                }
                break;
            case Tag::Object: {
                auto obj_type = type()->expect<ClassType>();
                auto obj_value = toObject();
                auto attributes = obj_type->getAttributes();
                for (const auto &attr : attributes) {
                    auto attribute = obj_value->getAttr(attr.getName());
                    attribute.
                    visit(visitor);
                }
                break;
            }
            case Tag::PyObject: {
                c10::intrusive_ptr<at::ivalue::PyObjectHolder> py_obj = toPyObjectHolder();
                auto match = py_obj->tryToInferType();
                if (match.success()) {
                    auto contained_value = py_obj->toIValue(match.type());
                    contained_value.
                    visit(visitor);
                }
                break;
            }
            default:
                break;
        }
    }

    void IValue::getSubValues(HashAliasedIValues &subValues) const {
        switch (this->tag) {
            case Tag::Tensor:
                subValues.insert(*this);
                return;
            case Tag::Tuple:
            case Tag::GenericList: {
                subValues.insert(*this);
                c10::ArrayRef<IValue> elems;
                if (isTuple()) {
                    elems = this->toTupleRef().elements();
                } else {
                    elems = this->toListRef();
                }
                for (auto &elem : elems) {
                    elem.getSubValues(subValues);
                }
                break;
            }
            case Tag::GenericDict:
                subValues.insert(*this);
                for (const auto &pair : this->toGenericDict()) {
                    pair.value().getSubValues(subValues);
                    pair.key().getSubValues(subValues);
                }
                break;
            case Tag::Object: {
                // Record Object IValue and its attributes.
                subValues.insert(*this);
                auto obj_type = type()->expect<ClassType>();
                auto obj_value = toObject();
                auto attributes = obj_type->getAttributes();
                for (const auto &attr: attributes) {
                    auto attribute = obj_value->getAttr(attr.getName());
                    attribute.getSubValues(subValues);
                }
                break;
            }
            case Tag::PyObject: {
                subValues.insert(*this);
                c10::intrusive_ptr<at::ivalue::PyObjectHolder> py_obj = toPyObjectHolder();
                auto match = py_obj->tryToInferType();
                TORCH_CHECK_TYPE(match.success(),
                                 "Cannot infer type of ", py_obj->toStr(), ": ", match.reason());
                auto contained_value = py_obj->toIValue(match.type());
                contained_value.getSubValues(subValues);
                break;
            }
            case Tag::Future:
            case Tag::Device:
            case Tag::Uninitialized:
            case Tag::Capsule:
                TORCH_CHECK_TYPE(
                        false, "Cannot inspect value of type ", this->tagKind());
                // Fall through
            default:
                // don't record scalars.
                break;
        }
    }

    bool IValue::overlaps(const IValue &rhs) const {
        HashAliasedIValues rhsSubValues, thisSubValues;
        rhs.getSubValues(rhsSubValues);
        getSubValues(thisSubValues);
        for (auto &sub : thisSubValues) {
            if (rhsSubValues.count(sub)) {
                return true;
            }
        }
        return false;
    }

    bool operator!=(const IValue &lhs, const IValue &rhs) {
        return !(lhs == rhs);
    }

    bool operator==(const IValue &lhs, const IValue &rhs) {
        IValue eq = lhs.equals(rhs);
        if (eq.isBool()) {
            return eq.toBool();
        }
        // The only case we don't return bool is for tensor comparison. In Python,
        // `bool()` is called on the return value of `__eq__` if the return value is
        // not a boolean. Mimic that behavior here.
        TORCH_INTERNAL_ASSERT(eq.isTensor());
        return eq.toTensor().is_nonzero();
    }

    bool IValue::ptrEqual(const IValue &lhs, const IValue &rhs) {
        TORCH_INTERNAL_ASSERT(lhs.is_intrusive_ptr);
        TORCH_INTERNAL_ASSERT(rhs.is_intrusive_ptr);
        return lhs.tag == rhs.tag &&
               lhs.payload.u.as_intrusive_ptr == rhs.payload.u.as_intrusive_ptr;
    }

    IValue IValue::equals(const IValue &rhs) const {
        const IValue &lhs = *this;
        switch (lhs.tag) {
            case Tag::None:
                // In Python you're not supposed to do this comparison apparently. Not
                // sure if we should warn here or what
                return rhs.isNone();
            case Tag::Tensor: {
                if (!rhs.isTensor()) {
                    return false;
                }
                return lhs.toTensor().eq(rhs.toTensor());
            }
            case Tag::Storage:
                return rhs.isStorage() && lhs.toStorage().unsafeGetStorageImpl() == rhs.toStorage().unsafeGetStorageImpl();
            case Tag::Double:
                return rhs.isDouble() && lhs.toDouble() == rhs.toDouble();
            case Tag::ComplexDouble:
                return rhs.isComplexDouble() && lhs.toComplexDouble() == rhs.toComplexDouble();
            case Tag::Int:
                return rhs.isInt() && lhs.toInt() == rhs.toInt();
            case Tag::Bool:
                return rhs.isBool() && lhs.toBool() == rhs.toBool();
            case Tag::String:
                return rhs.isString() && lhs.toStringRef() == rhs.toStringRef();
            case Tag::GenericDict:
                return rhs.isGenericDict() && lhs.toGenericDict() == rhs.toGenericDict();
            case Tag::Tuple:
                return rhs.isTuple() && *lhs.toTuple() == *rhs.toTuple();
            case Tag::Stream:
                return rhs.isStream() && lhs.toStream() == rhs.toStream();
            case Tag::Device:
                return rhs.isDevice() && lhs.toDevice() == rhs.toDevice();
            case Tag::GenericList:
                return rhs.isList() && lhs.toList() == rhs.toList();
            case Tag::Blob:
            case Tag::Future:
            case Tag::RRef:
            case Tag::Object:
            case Tag::PyObject:
            case Tag::Capsule:
            case Tag::Generator:
            case Tag::Quantizer:
                return ptrEqual(lhs, rhs);
            case Tag::Enum:
                return lhs.toEnumHolder()->is(*rhs.toEnumHolder());
            case Tag::Uninitialized:
                // Unitialized ivalues show up in no-ops when the compiler can prove a
                // value will never be used. Just return false on any equality comparison.
                return false;
        }
        // the above switch should be exhaustive
        TORCH_INTERNAL_ASSERT(false, "we should never reach here")
    }

    size_t IValue::hash(const IValue &v) {
        switch (v.tag) {
            case Tag::None:
                return 0;
            case Tag::Bool:
                return c10::get_hash(v.payload.u.as_bool);
            case Tag::Double:
                return c10::get_hash(v.payload.u.as_double);
            case Tag::Tensor:
                // Tensor __hash__ is equivalent to `id()`, so take the pointer value of
                // the tensor to emulate it
                return c10::get_hash(v.payload.as_tensor.unsafeGetTensorImpl());
                // NOLINTNEXTLINE(bugprone-branch-clone)
            case Tag::Storage:
                return c10::get_hash(v.payload.u.as_int);
            case Tag::Int:
                return c10::get_hash(v.payload.u.as_int);
            case Tag::String:
                return c10::get_hash(v.toStringRef());
            case Tag::Tuple:
                return c10::get_hash(*v.toTuple());
            case Tag::Device:
                return c10::get_hash(v.toDevice());
            case Tag::GenericDict:
            case Tag::GenericList:
            case Tag::Blob:
            case Tag::Future:
            case Tag::RRef:
            case Tag::Object:
            case Tag::PyObject:
            case Tag::Capsule:
            case Tag::Generator:
            case Tag::Quantizer:
            case Tag::ComplexDouble:
            case Tag::Enum:
            case Tag::Stream:
            case Tag::Uninitialized:
                throw std::runtime_error(
                        "unhashable type: '" + v.type()->repr_str() + "'");
        }
        // the above switch should be exhaustive
        TORCH_INTERNAL_ASSERT(false, "we should never reach here")
    }

    static bool isUndefinedTensor(const IValue &iv) {
        return iv.isTensor() && !iv.toTensor().defined();
    }

    bool IValue::is(const IValue &rhs) const {
        const IValue &lhs = *this;
        // Special handling for undefined tensors:
        // 1. Undefined_tensor is None and vice versa.
        if ((isUndefinedTensor(lhs) && rhs.isNone()) ||
            (lhs.isNone() && isUndefinedTensor(rhs))) {
            return true;
        }
        // 2. Undefined_tensor is Undefined_tensor.
        if (isUndefinedTensor(lhs) && isUndefinedTensor(rhs)) {
            return true;
        }

        if (lhs.isTensor()) {
            // Use the standard way of comparing two tensors for identity
            return rhs.isTensor() && lhs.toTensor().is_same(rhs.toTensor());
        }

        if (lhs.is_intrusive_ptr) {
            return rhs.is_intrusive_ptr && ptrEqual(lhs, rhs);
        }
        return lhs == rhs;
    }

    template<typename T>
    inline bool IValue::isListOf() const {
        // note: avoids calling type() to avoid extra referencing counting for the returned type.
        if (!isList()) {
            return false;
        }
        const auto &ty = static_cast<detail::ListImpl *>(payload.u.as_intrusive_ptr)->elementType;
        if (ty->kind() == T::Kind) {
            return true;
        }
        return *ty == *TypeFactory::get<T>();
    }

    bool IValue::isDoubleList() const {
        return isListOf<c10::FloatType>();
    }

    bool IValue::isComplexDoubleList() const {
        return isListOf<c10::ComplexType>();
    }

    bool IValue::isTensorList() const {
        return isListOf<c10::TensorType>();
    }

    bool IValue::isIntList() const {
        return isListOf<c10::IntType>();
    }

    bool IValue::isBoolList() const {
        return isListOf<c10::BoolType>();
    }
    IValue IValue::deepcopy() const {
        IValue::HashAliasedIValueMap memo;
        return deepcopy(memo);
    }

    IValue IValue::deepcopy(
            IValue::HashAliasedIValueMap& memo) const {
        if (memo.count(*this)) {
            return memo.at(*this);
        }
        IValue copy;
        switch(tag) {
            case IValue::Tag::Tensor:
                copy = IValue(toTensor().clone());
                break;
            case IValue::Tag::Tuple: {
                std::vector<IValue> copied_tuple;
                for (const auto& e : toTupleRef().elements()) {
                    copied_tuple.push_back(e.deepcopy(memo));
                }
                copy = IValue(ivalue::Tuple::create(copied_tuple));
            }
                break;
            case IValue::Tag::GenericList: {
                auto list = toList();
                auto copied_list = c10::impl::GenericList(list.elementType());
                for (IValue v : list) {
                    copied_list.push_back(v.deepcopy(memo));
                }
                copy = IValue(copied_list);
            }
                break;
            case IValue::Tag::GenericDict: {
                auto dict = toGenericDict();
                auto copied_dict = c10::impl::GenericDict(dict.keyType(), dict.valueType());
                for (const auto& entry : dict) {
                    copied_dict.insert(entry.key().deepcopy(memo), entry.value().deepcopy(memo));
                }
                copy = IValue(copied_dict);
            }
                break;
            case IValue::Tag::Object: {
                auto class_type = type()->expect<ClassType>();
                if (class_type->hasMethod("__getstate__") &&
                    class_type->hasMethod("__setstate__")) {
                    copy = ivalue::Object::create(
                            c10::StrongTypePtr(class_type->compilation_unit(), type()),
                            class_type->numAttributes());
                    auto state = class_type->getMethod("__getstate__")({*this});
                    class_type->getMethod("__setstate__")({copy, std::move(state)});
                } else {
                    copy = IValue(toObject()->deepcopy(memo));
                }
            } break;
            case IValue::Tag::Enum: {
                auto enum_holder = toEnumHolder();
                copy = IValue(c10::make_intrusive<ivalue::EnumHolder>(
                        enum_holder->type(),
                        enum_holder->name(),
                        enum_holder->value().deepcopy(memo)));
            } break;
            case IValue::Tag::String:
            case IValue::Tag::None:
            case IValue::Tag::Double:
            case IValue::Tag::Int:
            case IValue::Tag::Bool:
            case IValue::Tag::Device:
            case IValue::Tag::Uninitialized: {
                copy = *this;
            } break;
            default: {
                AT_ERROR("Can't deepcopy IValue with tag: ", tagKind());
            }
        }
        // NB: this doesn't work if an object contains itself, and it may
        // come up in the future when we expand the object system, we will
        // have a follow up PR to fix this when it becomes an issue.
        if (!isAliasOf(copy)) {
            memo[*this] = copy;
        }
        return copy;
    }

    void IValue::reportToTensorTypeError() const {
        TORCH_CHECK(false, "Expected Tensor but got ", tagKind());
    }

};